//===----------------------------------------------------------------------===//
//
//                         BusTub
//
// aggregation_executor.cpp
//
// Identification: src/execution/aggregation_executor.cpp
//
// Copyright (c) 2015-2021, Carnegie Mellon University Database Group
//
//===----------------------------------------------------------------------===//
#include <memory>
#include <vector>

#include "execution/executors/aggregation_executor.h"

namespace bustub {

AggregationExecutor::AggregationExecutor(ExecutorContext *exec_ctx, const AggregationPlanNode *plan,
                                         std::unique_ptr<AbstractExecutor> &&child)
    : AbstractExecutor(exec_ctx),
      plan_(plan),
      child_(std::move(child)),
      aht_(plan->GetAggregates(), plan_->GetAggregateTypes()),
      aht_iterator_(aht_.Begin()) {}

void AggregationExecutor::Init() {
  child_->Init();
  Tuple tuple{};
  RID rid{};
  while (child_->Next(&tuple, &rid)) {
    for (const auto &expr : plan_->GetGroupBys()) {
      const bustub::Value v = expr->Evaluate(&tuple, child_->GetOutputSchema());
      if (v.IsNull()) {
        // group by null value is not allowed..
        continue;
      }
    }
    aht_.InsertCombine(MakeAggregateKey(&tuple), MakeAggregateValue(&tuple));
    cnt_++;
  }
  aht_iterator_ = aht_.Begin();
}

auto AggregationExecutor::Next(Tuple *tuple, RID *rid) -> bool {
  if (cnt_ == 0) {
    if (!plan_->GetGroupBys().empty()) {
      return false;
    }
    std::vector<Value> values;
    auto aggregate_types = plan_->GetAggregateTypes();
    for (size_t i = 0; i < plan_->GetAggregates().size(); i++) {
      if (aggregate_types[i] == AggregationType::CountStarAggregate) {
        values.push_back(ValueFactory::GetIntegerValue(0));
      } else {
        values.push_back(ValueFactory::GetNullValueByType(TypeId::INTEGER));
      }
    }

    *tuple = Tuple(values, &GetOutputSchema());
    cnt_++;
    return true;
  }
  if (aht_iterator_ != aht_.End()) {
    std::vector<bustub::Value> vals{};
    for (const auto &group_by_val : aht_iterator_.Key().group_bys_) {
      vals.push_back(group_by_val.GetAs<bustub::Value>());
    }
    for (const auto &aggregate_value : aht_iterator_.Val().aggregates_) {
      vals.push_back(aggregate_value.GetAs<bustub::Value>());
    }
    Tuple result{vals, &plan_->OutputSchema()};
    *tuple = result;
    ++aht_iterator_;
    return true;
  }
  return false;
}

auto AggregationExecutor::GetChildExecutor() const -> const AbstractExecutor * { return child_.get(); }

}  // namespace bustub
